import tushare as ts
from os import mkdir, path, getcwd
from stock_db import StockOB

db = StockOB()


class TushareStock:
    def __init__(self):
        pass

    def get_df(self, start_date, end_date, code):
        """
        获取股票数据dataframe数据
        """
        try:
            # 从数据库获取token
            get_token_result = db.select_token()
            # 获取token成功
            if get_token_result["status"] == 200:
                ts_token = get_token_result["ts_token"]
                ts.set_token(ts_token)
                # print(ts_token)
            # 获取token失败
            else:
                # 返回sql执行的报错信息
                return get_token_result
            # 用token查询股票数据
            pro = ts.pro_api()
            # 添加后缀查询数据
            ends = [".SH", ".SZ", ".BJ", ".HK"]
            for e in ends:
                df = pro.daily(ts_code=code+e, start_date=start_date, end_date=end_date).sort_values(by="trade_date")
                if not df.empty:
                    # 查询数据返回空数据
                    if not df.empty:
                        return {
                            "status": 200,
                            "msg": f"获取【{code}】股票数据成功！",
                            "df": df
                        }
            # 没有查询到数据
            return {
                "status": 400,
                "ts_code": code,
                "msg": f"获取【{code}】股票数据失败！"
            }
        # 代码报错
        except Exception as e:
            # token填写错误
            result = {
                "status": 400,
                "msg": str(e)
            }
            return result

    def get_stock(self, start_date, end_date, code):
        # 获取股票数据
        result = self.get_df(start_date, end_date, code)
        # 获取股票数据失败
        if result["status"] != 200:
            return result
        # 获取股票数据成功
        df = result["df"]
        # 预测信息【后续加上？？？？？】
        forecast_msg = '暂无预测模型！'
        # 重命名
        df2 = df.rename(
            columns={'ts_code': "股票代码", 'trade_date': "交易日期", 'open': "开盘价", 'high': "最高价",
                     'low': "最低价",
                     'close': "收盘价", 'pre_close': "昨收价(前复权)",
                     'change': "涨跌额", 'pct_chg': "涨跌幅", 'vol': "成交量 （手）", 'amount': "成交额 （千元）"}
        )
        # 返回数据
        result = {
            "status": 200,
            "msg": f"获取【{code}】股票数据成功!",
            "stock_data": {
                "stock_trade_days": df2["交易日期"].values.tolist(),
                "stock_trade_data": df2[["开盘价", "收盘价", "最低价", "最高价"]].values.tolist(),
                "MA5": df2["收盘价"].rolling(4).mean().map(lambda x: round(x, 3)).fillna("-").tolist(),
                "MA10": df2["收盘价"].rolling(9).mean().map(lambda x: round(x, 3)).fillna("-").tolist(),
                "MA20": df2["收盘价"].rolling(19).mean().map(lambda x: round(x, 3)).fillna("-").tolist(),
                "MA30": df2["收盘价"].rolling(29).mean().map(lambda x: round(x, 3)).fillna("-").tolist(),
                "forecast": forecast_msg,
                "table_columns": df2.columns.tolist(),
                "table_data": df2.to_dict(orient="records"),
                "ts_name": code
            }
        }
        return result

    def download_stock(self, start_date, end_date, code):
        """
        下载【保存】数据
        """
        # 获取股票数据
        result = self.get_df(start_date, end_date, code)
        if result["status"] != 200:
            return result
        df = result["df"]
        ts_name = result["ts_name"]
        df2 = df.rename(
            columns={'ts_code': "股票代码", 'trade_date': "交易日期", 'open': "开盘价", 'high': "最高价",
                     'low': "最低价",
                     'close': "收盘价", 'pre_close': "昨收价(前复权)",
                     'change': "涨跌额", 'pct_chg': "涨跌幅", 'vol': "成交量 （手）", 'amount': "成交额 （千元）"}
        )
        # 保存csv数据的名称
        file_name = f"{ts_name}_{code}_{start_date}_{end_date}_数据.csv"
        # 保存数据的地址
        file_path = path.join(path.abspath(path.dirname(getcwd())), "股票数据")
        # 如果不存在就创建
        if not path.exists(file_path):
            mkdir(file_path)
        csv_file_path = path.join(file_path, file_name)
        # 保存数据
        df2.to_csv(csv_file_path, index=False)
        result = {
            "status": 200,
            "csv_file_path": csv_file_path,
            "msg": f"数据已保存在本地【{file_path}】文件夹中！"
        }
        return result
